#include <Windows.h>
#include <osg/geode>
#include <osgDB/Registry>
#include <osgDB/WriteFile>
#include <osg/Notify>
#include <osg/geometry>
#include <osg/PositionAttitudeTransform>
#include <osg/material>
#include <osg/texture2d>
#include <osg/image>
#include <osg/AlphaFunc>
#include <osg/BlendFunc>
#include <osg/texmat>

#include <inode.h>
#include <bitmap.h>

#include "maxplugin/osg_scene_saver.h"
#include "maxplugin/model.h"
#include "maxplugin/texture.h"
#include "maxplugin/dialog.h"
#include "maxplugin/file_system.h"

osg_scene_saver::osg_scene_saver(TimeValue time, bool export_selected)
	:	scene_saver(time, export_selected)
{ }

osg_scene_saver::~osg_scene_saver() { }

bool osg_scene_saver::save(const char *path)
{
	file_system &fs = file_system::get_singleton();
	osg::ref_ptr<osg::PositionAttitudeTransform> root = new osg::PositionAttitudeTransform;
	osg_mesh_map meshes = convert_meshes();
	osg_mtl_map materials = convert_materials();
	for(int i = 0; i < get_num_models(); i++) {
		model_ptr m = get_model(i);
		osg_node_ptr node = convert_model(m, meshes, materials);

		if(dialog::get_singleton().scene_export_option() == SO_EXPORT_SCENE)
			root->addChild(node);
		else {
			osg::ref_ptr<osg::PositionAttitudeTransform> scene = new osg::PositionAttitudeTransform;
			scene->addChild(node);
			osgDB::writeNodeFile(*scene.get(), m->get_name() + fs.parse_suffix(path));
		}
	}
	if(dialog::get_singleton().scene_export_option() != SO_EXPORT_SCENE)
		return true;
	return osgDB::writeNodeFile(*root.get(), fs.parse_filename(path));
}

osg_scene_saver::osg_node_ptr osg_scene_saver::convert_model(model_ptr model,
		const osg_mesh_map &meshes, const osg_mtl_map &materials)
{
	osg::ref_ptr<osg::PositionAttitudeTransform> node = new osg::PositionAttitudeTransform;
	node->setName(model->get_name());

	vector3 position, scale;
	quaternion orientation;
	model->get_transforms().decomposition(position, scale, orientation);
	osg::Quat rot(orientation.x, orientation.y, orientation.z, orientation.w);
	osg::Vec3 pos(position.x, position.y, position.z);
	osg::Vec3 sc(scale.x, scale.y, scale.z);
	node->setAttitude(rot);
	node->setPosition(pos);
	node->setScale(sc);

	for(int j = 0; j < model->get_num_meshes(); j++) {
		mesh_ptr mesh = model->get_mesh(j);
		osg_mesh_map::const_iterator it = meshes.find(mesh->get_name());
		osg::ref_ptr<osg::Geode> geode = new osg::Geode;
		geode->setDataVariance(osg::Object::STATIC);
		geode->addDrawable(it->second.get());
		geode->setName(mesh->get_name());

		if(model->get_num_materials() > j) {
			material_ptr mat = model->get_material(j);
			osg_mtl_map::const_iterator it = materials.find(mat->get_name());
			geode->setStateSet(it->second);
		}

		node->addChild(geode);
	}

	for(int i = 0; i < model->get_num_childrens(); i++) {
		model_ptr m = model->get_child(i);
		osg_node_ptr child = convert_model(m, meshes, materials);
		node->addChild(child);
	}

	return node;
}

osg_scene_saver::osg_mesh_map osg_scene_saver::convert_meshes()
{
	osg_mesh_map meshes;
	for(int i = 0; i < get_num_meshes(); i++) {
		mesh_ptr m = get_mesh(i);
		osg::ref_ptr<osg::Geometry> geom = new osg::Geometry;

		osg::ref_ptr<osg::Vec3Array> verts = new osg::Vec3Array;
		geom->setVertexArray(verts);
		for(int j = 0; j < m->get_num_vertices(); j++) {
			vector3 v = m->get_coord(j);
			verts->push_back(osg::Vec3(v.x, v.y, v.z));
		}

		osg::ref_ptr<osg::UShortArray> indices = new osg::UShortArray;
		geom->setVertexIndices(indices);
		geom->setTexCoordIndices(0, indices);
		for(int j = 0; j < m->get_num_indices(); j++)
			indices->push_back(m->get_idx(j));

		osg::ref_ptr<osg::Vec3Array> norms = new osg::Vec3Array;
		geom->setNormalArray(norms);
		for(int j = 0; j < m->get_num_vertices(); j++) {
			vector3 n = m->get_normal(j);
			norms->push_back(osg::Vec3(n.x, n.y, n.z));
		}
		geom->setNormalBinding(osg::Geometry::BIND_PER_VERTEX);
		geom->setNormalIndices(indices.get());

		osg::ref_ptr<osg::Vec2Array> texc = new osg::Vec2Array;
		geom->setTexCoordArray(0, texc);
		for(int j = 0; j < m->get_num_vertices(); j++) {
			vector2 t = m->get_texcoord(j);
			texc->push_back(osg::Vec2(t.x, t.y));
		}

		geom->addPrimitiveSet(new osg::DrawArrays(osg::PrimitiveSet::TRIANGLES, 0, indices->size()));
		geom->setName(m->get_name());
		meshes[m->get_name()] = geom;
	}
	return meshes;
}

osg_scene_saver::osg_mtl_map osg_scene_saver::convert_materials()
{
	osg_mtl_map materials;
	osg_tex_map textures = convert_textures();
	for(int i = 0; i < get_num_materials(); i++) {
		material_ptr stdmat = get_material(i);
		osg::ref_ptr<osg::Material> osgmat = new osg::Material();

		vector4 a = stdmat->get_ambient_colour();
		vector4 d = stdmat->get_diffuse_colour();
		vector4 s = stdmat->get_specular_colour();

		osgmat->setAmbient(osg::Material::FRONT_AND_BACK, osg::Vec4(a[0], a[1], a[2], a[3]));
		osgmat->setDiffuse(osg::Material::FRONT_AND_BACK, osg::Vec4(d[0], d[1], d[2], d[3]));
		osgmat->setSpecular(osg::Material::FRONT_AND_BACK, osg::Vec4(s[0], s[1], s[2], s[3]));

		osg::ref_ptr<osg::StateSet> stateset = new osg::StateSet;
		stateset->setMode(GL_NORMALIZE, osg::StateAttribute::ON);
		if(stdmat->is_doublesided())
			stateset->setMode(GL_CULL_FACE,osg::StateAttribute::OFF);
		else
			stateset->setMode(GL_CULL_FACE,osg::StateAttribute::ON);

		texture_ptr tex = stdmat->get_diffuse();
		if(!tex.is_null()) {
			matrix4 uvtm = tex->get_uv_transforms().transpose();
			osg::Matrixf tm(uvtm[0]);
			osg::ref_ptr<osg::TexMat> texmat = new osg::TexMat(tm);
			stateset->setAttribute(texmat);
		}
		else
			stateset->setAttribute(osgmat.get());

		std::string diffuse;
		texture_ptr diffmap = stdmat->get_diffuse();
		if(!diffmap.is_null())
			diffuse = diffmap->get_name();
		osg_tex_map::iterator it = textures.find(diffuse);
		if(it != textures.end()) {
			stateset->setTextureAttributeAndModes(0, it->second.get(), osg::StateAttribute::ON);

			if(stdmat->has_transparency()) {
				stateset->setMode(GL_BLEND,osg::StateAttribute::ON);
				stateset->setRenderingHint( osg::StateSet::TRANSPARENT_BIN );
			}
		}

		materials[stdmat->get_name()] = stateset;
	}
	return materials;
}

osg_scene_saver::osg_tex_map osg_scene_saver::convert_textures()
{
	osg_tex_map textures;
	for(int i = 0; i < get_num_textures(); i++) {
		texture_ptr t = get_texture(i);
		t->save();
		osg::ref_ptr<osg::Image> img = new osg::Image;
		img->setFileName(t->get_name());
		osg::ref_ptr<osg::Texture2D> tex = new osg::Texture2D();
		tex->setImage(img.get());
		tex->setInternalFormatMode(osg::Texture::USE_IMAGE_DATA_FORMAT);
		tex->setWrap(osg::Texture::WRAP_S, osg::Texture::REPEAT);
		tex->setWrap(osg::Texture::WRAP_T, osg::Texture::REPEAT);
		tex->setWrap(osg::Texture::WRAP_R, osg::Texture::REPEAT);
		textures[t->get_name()] = tex;
	}
	return textures;
}

void osg_scene_saver::write_skeletons() { }